import os
from modelscope.trainers import build_trainer
from modelscope.msdatasets import MsDataset
from modelscope.utils.constant import DownloadMode
from modelscope.trainers.hooks import Hook, Priority

# 3. 训练配置
def train_model():
    model_id = 'damo/nlp_structbert_siamese-uninlu_chinese-base'
    WORK_DIR = 'models'

    # 加载数据集
    train_dataset = MsDataset.load(
        'people_daily_ner_1998_tiny',
        namespace='damo',
        split='train',
        download_mode=DownloadMode.FORCE_REDOWNLOAD
    )
    eval_dataset = MsDataset.load(
        'people_daily_ner_1998_tiny',
        namespace='damo',
        split='validation',
        download_mode=DownloadMode.FORCE_REDOWNLOAD
    )

    # 关键配置：通过 custom_hooks 添加回调
    kwargs = {
        'model': model_id,
        'model_revision':'master',
        'train_dataset': train_dataset,
        'eval_dataset': eval_dataset,
        'max_epochs': 3,
        'work_dir': WORK_DIR
    }

    # 构建训练器
    trainer = build_trainer('siamese-uie-trainer', default_args=kwargs)
    trainer.train()


if __name__ == '__main__':
    train_model()